package datastructure.tree;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

/**
 *  Create a linked list of all the nodes at each depth
 * 
 * 
 * @author Lei
 */
public class GetNodeListForEachLevel {

    public static List<List<TreeNode>> getNodeListForEachLevel(TreeNode root) {

        if (root == null) {
            return null;
        }

        List<List<TreeNode>> retList = new ArrayList<List<TreeNode>>();
        int currentLevelNodeCount = 1;
        int nextLevelNodeCount = 0;

        Queue<TreeNode> queue = new LinkedList<TreeNode>();
        List<TreeNode> aList = new ArrayList<TreeNode>();
        queue.add(root);

        while (!queue.isEmpty()) {
            TreeNode n = queue.poll();
            currentLevelNodeCount--;
            if (n.left != null) {
                queue.add(n.left);
                nextLevelNodeCount++;
            }

            if (n.right != null) {
                queue.add(n.right);
                nextLevelNodeCount++;
            }

            aList.add(n);
            if (currentLevelNodeCount == 0) {
                List<TreeNode> aNewList = new ArrayList<TreeNode>();
                for (TreeNode aNode : aList) {
                    aNewList.add(aNode);
                }
                retList.add(aNewList);
                aList.clear();

                currentLevelNodeCount = nextLevelNodeCount;
                nextLevelNodeCount = 0;
            }
        }

        return retList;
    }

    public static void main(String[] args) {

        TreeNode n1 = new TreeNode(1);
        TreeNode n2 = new TreeNode(2);
        TreeNode n3 = new TreeNode(3);
        TreeNode n4 = new TreeNode(4);
        TreeNode n5 = new TreeNode(5);
        TreeNode n6 = new TreeNode(6);
        TreeNode n7 = new TreeNode(7);
        
        n5.left = n2;
        n5.right = n6;
        n6.left = n7;

        n2.left = n1;
        n2.right = n4;

        n4.left = n3;

        List<List<TreeNode>> retList = GetNodeListForEachLevel.getNodeListForEachLevel(n5);
        for (List<TreeNode> aList : retList) {
            System.out.println("a LIST: ");
            for (TreeNode aNode : aList) {
                System.out.println(aNode.value);
            }

            System.out.println("   ");
        }

    }
}
